import os
import sys
import time, datetime
import pandas as pd
import requests
from selenium import webdriver
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.common.by import By
from selenium.webdriver.support import expected_conditions as EC
import json
from io import BytesIO
from src.exts import TestLogger
from src.mxtf_api import MxtfApi
from src.any2txt import Any2Txt

class ExchangeFileSpider:
    def __init__(self,model='windows',updatedb=False,logger=None):
        if logger:
            self.logger = logger
        else:
            self.logger = TestLogger()
        self.updatedb = updatedb # 是否写入数据库，默认不写入，打印出结果
        self.summary = {'measure_index':None, # 监管发函索引
                        'file_urls':None} # 监管发函url
        self.model = model
        
        self.chromedriver_dir = 'D:/Business/财富管理投顾产品/scripts/' #chromedriver位置
        
        #self.tmp_file_dir = 'C:/Users/Julin/Downloads/' #临时文件暂存位置，暂未使用
        self.exchange_file_dir = 'C:/Users/Julin/Downloads/' #监管文件存放位置，测试用
        
        self.headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"}
        
        self.sz_exchange_domain = 'https://www.szse.cn/'
        
        self.sz_exchange_jgcs_file_index_url = 'https://www.szse.cn/disclosure/supervision/measure/measure/index.html' #深交所监管措施列表
        self.sz_exchange_jgcs_file_index_download_btn_class = 'btn-default-excel btn-default-hasicon' #下载按钮class name
        self.sz_exchange_jgcs_file_index_download_api = '/api/report/ShowReport?SHOWTYPE=xlsx&CATALOGID=1800_jgxxgk&TABKEY=tab1&selectBkmc=0'
        self.sz_exchange_jgcs_file_index_header = ['stock_code','corp_name','measure_date','measure_subtype','measure_file_name','measure_target']
        self.sz_exchange_jgcs_file_url_prefix = 'https://reportdocs.static.szse.cn/UpFiles/zqjghj/' #深交所监管措施，2024年之后
        
        self.sz_exchange_jlcf_file_index_url = 'https://www.szse.cn/disclosure/supervision/measure/pushish/index.html' #深交所纪律处分列表
        self.sz_exchange_jlcf_file_index_download_api = '/api/report/ShowReport?SHOWTYPE=xlsx&CATALOGID=1800_jgxxgk_cf&TABKEY=tab2&selectGsbk=0'
        self.sz_exchange_jlcf_file_index_download_btn_class = 'btn-default-excel btn-default-hasicon' #下载按钮class name
        self.sz_exchange_jlcf_file_index_header = ['stock_code','corp_name','measure_date','measure_subtype','measure_title','measure_file_name']
        self.sz_exchange_jlcf_file_url_prefix = 'https://reportdocs.static.szse.cn/UpFiles/zqjghj/' #深交所纪律处分，2024年之后  

        self.sh_exchange_domain = 'https://www.sse.com.cn/'
        self.sh_exchange_jgcs_file_index_url = 'https://www.sse.com.cn/regulation/supervision/measures/' #上交所纪律处分列表
        self.sh_exchange_file_index_header = ['stock_code','corp_name','measure_subtype','measure_title','measure_target','measure_date','measure_file_url']
        
        self.exchange_file_index = None
        self.exchange_file_urls = None
        self.session = requests.Session()
        self.local_file_seq = int(datetime.datetime.now().strftime("%Y%m%d%H%M%S%f"))
        
        self.mxtfapi = MxtfApi()
        self.any2txt = Any2Txt()

    def tmp_file_path(self,filetype): # 暂未使用
        seqstamp = time.strftime("%Y%m%d%H%M%S",time.localtime())
        tmp_file_path = self.tmp_file_dir + 'tmp_' + seqstamp + '.' + filetype
        return tmp_file_path

    def local_file_name(self,filename):
        filename = filename.split('.')[0] + '__' + str(self.local_file_seq) + '.' + filename.split('.')[-1]
        self.local_file_seq += 1
        return filename

    @staticmethod
    def df2json(df):
        return df.to_json(orient='records')
    
    def add2db(self,table_name,data_json):
        if self.updatedb:
            results = self.mxtfapi.execute_api(table_name,'insert',data_json)
        else:
            results = ''
        return results
    
    def get_exchange_file_index(self):
        self.logger.info(f"获取上海交易所监管文件索引")
        sh_exchange_file_index = self.get_sh_exchange_file_index()
        self.logger.info(f"获取上海交易所监管文件索引完成")
        self.logger.info(f"获取深圳交易所监管文件索引")
        sz_exchange_file_index = self.get_sz_exchange_file_index()
        self.logger.info(f"获取深圳交易所监管文件索引完成")
        exchange_file_index = pd.concat([sh_exchange_file_index,sz_exchange_file_index], ignore_index=True)
        exchange_file_index['measure_file_local_name'] = [self.local_file_name(filename) for filename in list(exchange_file_index['measure_file_name'])]         
        self.exchange_file_index = exchange_file_index
    
    def get_sh_exchange_file_index(self):
        if self.model=='windows':
            os.environ["webdriver.ie.driver"] = self.chromedriver_dir
        option=webdriver.ChromeOptions()
        option.add_argument('--headless')
        option.add_argument("--no-sandbox")
        option.add_argument("--disable-dev-shm-usage")

        driver = webdriver.Chrome(options=option)
        
        driver.implicitly_wait(5) 
        driver.get(self.sh_exchange_jgcs_file_index_url)
        
        driver.set_window_size(1920, 1080)
        
        # element = driver.find_element_by_xpath("//*[@class='btn dropdown-toggle btn-light']")
        element = WebDriverWait(driver, 10).until(
            EC.visibility_of_element_located((By.XPATH, "//*[@class='btn dropdown-toggle btn-light']"))
        )
        driver.execute_script("arguments[0].scrollIntoView();", element)
        element.click()
        driver.find_element_by_xpath("//*[text()='每页100条']").click()
        time.sleep(5)
        
        pgcount = driver.find_element_by_xpath("//*[text()='下一页']/../preceding-sibling::*[1]").text
        pgcount = int(pgcount)
        fultextlist = []
        try:
            for n in range(pgcount):
                pgtabtrlist = driver.find_elements_by_xpath("//*[@class = 'table table-hover ']/tbody/tr")
                for eachtr in pgtabtrlist:
                    pgtabtdlist = eachtr.find_elements_by_xpath(".//td")
                    currtextlist = []
                    for eachtd in pgtabtdlist:
                        currtextlist += [eachtd.text]
                    try:
                        curr_href = eachtr.find_element_by_xpath(".//*[@class = 'table_td_xlwrap']").get_attribute('href')
                    except:
                        curr_href = ''
                    currtextlist += [str(curr_href)]
                    fultextlist += [currtextlist]
                self.logger.debug(f"爬取上海交易所监管文件索引第{n}页")
                if n < pgcount-1:
                    driver.find_element_by_xpath("//*[text()='下一页']").click()
                    time.sleep(5)
        except Exception as e:
            self.logger.error(f"爬取上海交易所监管文件索引失败:{e}")
        self.logger.debug("爬取上海交易所监管文件索引完成")
        sh_exchange_file_index = pd.DataFrame(fultextlist)
        sh_exchange_file_index.columns = self.sh_exchange_file_index_header
        sh_exchange_file_index['measure_file_name'] = sh_exchange_file_index['measure_file_url'].apply(lambda x: x.split('/')[-1])
        sh_exchange_file_index['measure_entity'] = 'sh'
        sh_exchange_file_index['measure_type'] = 'jgcs'
        
        driver.close()
        
        return sh_exchange_file_index

    def get_sz_exchange_file_index(self):
        
        sz_jgcs_file_index = self.get_sz_exchange_jgcs_file_index()
        sz_jlcf_file_index = self.get_sz_exchange_jlcf_file_index()
        sz_exchange_file_index = pd.concat([sz_jgcs_file_index,sz_jlcf_file_index], ignore_index=True)

        return sz_exchange_file_index

    def get_sz_exchange_jgcs_file_index(self):
        
        sz_jgcs_api_url = self.sz_exchange_domain + self.sz_exchange_jgcs_file_index_download_api
        sz_jgcs_file_index = self.get_sz_exchange_file_index_by_api(sz_jgcs_api_url)
        sz_jgcs_file_index.columns = self.sz_exchange_jgcs_file_index_header
        self.logger.debug(f"爬取深圳交易所监管文件{sz_jgcs_api_url} ，数量{sz_jgcs_file_index.count()}")
        sz_jgcs_file_index['measure_entity'] = 'sz'
        sz_jgcs_file_index['measure_type'] = 'jgcs'
        sz_jgcs_file_index['measure_file_url'] = sz_jgcs_file_index['measure_file_name'].apply(lambda x: self.sz_exchange_jgcs_file_url_prefix + x)
        
        return sz_jgcs_file_index
    
    def get_sz_exchange_jlcf_file_index(self):
        
        sz_jlcf_api_url = self.sz_exchange_domain + self.sz_exchange_jlcf_file_index_download_api
        sz_jlcf_file_index = self.get_sz_exchange_file_index_by_api(sz_jlcf_api_url)
        sz_jlcf_file_index.columns = self.sz_exchange_jlcf_file_index_header
        self.logger.debug(f"爬取深圳交易所监管文件{sz_jlcf_api_url} ，数量{sz_jlcf_file_index.count()}")
        sz_jlcf_file_index['measure_entity'] = 'sz'
        # sz_jlcf_file_index['measure_type'] = 'jgcs'
        sz_jlcf_file_index['measure_type'] = 'jlcf'
        sz_jlcf_file_index['measure_file_url'] = sz_jlcf_file_index['measure_file_name'].apply(lambda x: self.sz_exchange_jlcf_file_url_prefix + x)
        
        return sz_jlcf_file_index

    def get_sz_exchange_file_index_by_api(self,api_url):
        
        response = requests.get(api_url, headers=self.headers)

        if response.status_code == 200:
            #new_tmp_file_path = self.tmp_file_path('xlsx')
            #with open(new_tmp_file_path, 'wb') as file:
                #file.write(response.content)
        
            #data_df = pd.read_excel(new_tmp_file_path, dtype=str)
            
            # 使用BytesIO处理二进制流
            bytes_data = BytesIO(response.content)
            # 使用Pandas读取Excel文件
            data_df = pd.read_excel(bytes_data, engine='openpyxl')  # or 'xlrd' for xls
        
        return data_df

    def get_exchange_file_urls(self,st_date,ed_date):
        
        self.exchange_file_index['f_measure_date'] = pd.to_numeric(self.exchange_file_index['measure_date'].str.replace('-',''))
        self.exchange_file_urls = self.exchange_file_index[(self.exchange_file_index['measure_file_url'].str.contains('http')) & \
                                                     (self.exchange_file_index['f_measure_date']>=st_date) & \
                                                     (self.exchange_file_index['f_measure_date']<=ed_date)]\
                                                    [['measure_file_url','measure_file_local_name']]        
        self.exchange_file_index = self.exchange_file_index.drop(['f_measure_date'], axis=1)
        
    def download_exchange_file(self):
        
        for fileurl, filename in self.exchange_file_urls.values:
            try:
                response = self.session.get(fileurl)
                if response.status_code == 200:
                    with open(self.exchange_file_dir+filename, 'wb') as f:
                        f.write(response.content)
                        time.sleep(1)
            except:
                pass
            
    def run_all(self,st_date,ed_date):
        try:
            self.logger.info(f'start to 监管数据')
            self.get_exchange_file_index()
        except Exception as e:
            self.logger.error(f'get_exchange_file_index fail {e}')
            return None
            
        self.get_exchange_file_urls(st_date,ed_date)
        
        if self.updatedb:
            self.db_response = self.add2db('exchange_file_index',self.df2json(self.exchange_file_index))
            self.download_exchange_file(st_date,ed_date)
        else:
            self.summary['measure_index'] = self.exchange_file_index.to_json(orient='records')
            self.summary['file_urls'] = self.exchange_file_urls.to_json(orient='records')
            summary_json = json.dumps(self.summary)
            return summary_json
    def get_exchange_file(self,exchange):
        result = None
        if exchange == 'sh':
            self.logger.info(f"获取上海交易所监管文件索引")
            result = self.get_sh_exchange_file_index()
            self.logger.info(f"获取上海交易所监管文件索引完成")
        else:
            self.logger.info(f"获取深圳交易所监管文件索引")
            result = self.get_sz_exchange_file_index()
            self.logger.info(f"获取深圳交易所监管文件索引完成")
        
        # 删除股票代码为空的数据
        result = result.dropna(subset=['stock_code'])
        result = result.dropna(subset=['measure_date'])
        result['measure_file_local_name'] = [self.local_file_name(filename) for filename in list(result['measure_file_name'])]
        # 股票代码的数字格式转为6位代码字符
        # result['stock_code'] = [ ('000000' + str(int(eachstockcode)))[-6:] for eachstockcode in list(result['stock_code'] if eachstockcode)]
        # 生成标题
        result['combined_title'] = result.apply(self.generate_title, axis=1)
        
        return result.to_json(orient='records')
    
    
    def generate_title(self,row):
        exchangecode2name = {'sh':'上交所','sz':'深交所'}
        if row['measure_subtype'] == '监管函' or  row['measure_subtype'] == '监管工作函':
            return ('【' + row['measure_date'] + '】' + exchangecode2name[row['measure_entity']] +
                    '关于' + row['corp_name'] +
                    '(' + ('000000' + str(int(row['stock_code'])))[-6:] + ')' +
                    '的' + row['measure_subtype'])
        else:
            return '【' + row['measure_date'] + '】' + exchangecode2name[row['measure_entity']] + str(row['measure_title']).split('-')[-1]
        
    def process_exchange_file_data(self,exchange_file_data):
        
        exchange_file_data['measureContent'] = ''
        exchange_file_data['cosUrl'] = ''
        
        try:
            fileurl = exchange_file_data['measureFileUrl']
            response = self.session.get(fileurl)
            if response.status_code == 200:
                try:
                    up_response = self.mxtfapi.uploadFile(response.content)
                    exchange_file_data['cosUrl'] = up_response.get('data',None).get('url',None)
                except Exception as e:
                    self.logger.error(f"上传监管文件失败, url: {fileurl}; error: {e}")
                try:
                    filecontent = BytesIO(response.content)
                    filetextcontent = self.any2txt.pdf2txt(filecontent)
                    exchange_file_data['measureContent'] = filetextcontent
                except Exception as e:
                    self.logger.error(f"解析监管文件失败, url: {fileurl}; error: {e}")
            else:
                self.logger.error(f"获取监管文件失败, url: {fileurl}")
        except Exception as e:
            self.logger.error(f"获取监管文件失败, url: {fileurl}; error: {e}")
        
        return exchange_file_data
        
def main():
    
    if len(sys.argv) < 2:
        print("Usage: python ExchangeFileSpider.py <dates_json>")
        sys.exit(1)
        
    dates_json = sys.argv[1] #{"stDate":20241101,"edDate":20241231}
    st_date, ed_date = json.loads(dates_json).values()
    exchange_file_spider = ExchangeFileSpider(updatedb=False)
    exchange_file_spider.run_all(st_date,ed_date)
        
if __name__ == "__main__":
    main()
